"""
  此文件用来训练 头车自由驾驶的lstm模型
"""

# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)


# 导入外部包
import numpy as np
from sklearn.model_selection import train_test_split
from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout, Reshape

# 导入内部包
from utils.read_data import read_extract_free_drive_data
from utils.lstm import create_lstm_dataset, normalization

EPOCHS = 50 # 迭代次数
BATCH_SIZE = 64 # 一次迭代中的样本数量

FILE_PATH_I80_1_to = "../../Data/Ngsim数据集/I80数据集/3. 提取数据/3. 自由驾驶数据/trajectories-0400-0415_free_drive.txt" # 2896 150 7
FILE_PATH_I80_2_to = "../../Data/Ngsim数据集/I80数据集/3. 提取数据/3. 自由驾驶数据/trajectories-0500-0515_free_drive.txt" # 2697 150 7

def build_model(input_shape, output_shape):
  model = Sequential([
    LSTM(units=64, input_shape=input_shape, activation="relu", return_sequences=True, dropout=0.2, recurrent_dropout=0.2),
    Dropout(0.3),  # 添加独立Dropout层（训练和预测时均需激活）
    LSTM(units=32, activation="relu", dropout=0.2, recurrent_dropout=0.2),
    Dropout(0.2), # 普通Dropout层（默认仅在训练时激活）
    Dense(output_shape[-1] * output_shape[0], activation="linear"),
    Reshape(output_shape)
  ])
  model.compile(loss="mse", optimizer="adam")
  return model

if __name__ == "__main__":
  # 读取数据
  free_drive_datas = np.array(read_extract_free_drive_data(FILE_PATH_I80_1_to))

  # 归一化
  free_drive_datas, min_max_list = normalization(free_drive_datas, [4, 5], [True, False])
  print("速度的min和max: ", min_max_list[0][0], min_max_list[0][1]) # 0.74507 24.00588
  print("速加度的min和max: ", min_max_list[1][0], min_max_list[1][1]) # -2.50226 2.25379
  print("数据归一化后: ", free_drive_datas[0][0][4], free_drive_datas[0][0][5]) # 0.1317636832079364 -0.03555261193637571

  # 创建训练集和验证集
  n_past, n_future = 30, 10
  X, Y = create_lstm_dataset(free_drive_datas, n_past, n_future, [4, 5], [5]) # !历史是速度和加速度，预测是加速度
  x_train, x_val, y_train, y_val = train_test_split(X, Y, test_size=0.3, random_state=42)
  print("训练集数据: ", len(x_train), len(x_train[0]), len(x_train[0][0])) # 225019 30 2
  print("训练集标签", len(y_train), len(y_train[0]), len(y_train[0][0])) # 225019 10 1
  print("验证集数据: ", len(x_val), len(x_val[0]), len(x_val[0][0])) # 96437 30 2
  print("验证集标签", len(y_val), len(y_val[0]), len(y_val[0][0])) # 96437 10 1

  # 构建模型
  input_shape = (n_past, 2)
  output_shape = (n_future, 1)
  model = build_model(input_shape, output_shape)

  # 训练模型
  model.fit(x_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_data=(x_val, y_val), verbose=1)
  model.save("./model/free_drive_EPOCHS_50_BATCH_SIZE_64_unit_64_32_implement-normalization.keras")